import requests
import base64
import os
import hashlib


class FileUploader:
    def __init__(self, host, port, max_upload_size, chunk_size):
        self.host = host
        self.port = port
        self.max_upload_size = max_upload_size
        self.chunk_size = chunk_size
        self.url = f'http://{self.host}:{self.port}/file/post'

    @staticmethod
    def calculate_file_hash(data, hash_algorithm="sha1"):
        hash_obj = hashlib.new(hash_algorithm)
        if isinstance(data, str):
            data = data.encode('utf-8')
        hash_obj.update(data)
        return hash_obj.hexdigest()

    def upload_file(self, file_path, qr_save_dir="temp"):
        if not os.path.exists(qr_save_dir):
            os.makedirs(qr_save_dir)

        try:
            file_size = os.path.getsize(file_path)

            if file_size > self.max_upload_size:
                print(f'文件大小 {file_size} 超过阈值，启动分块上传')
                total_chunks = (file_size + self.chunk_size - 1) // self.chunk_size
                file_hash = self.calculate_file_hash(open(file_path, 'rb').read(), "sha1")

                # 初始化分块上传
                init_response = requests.post(
                    f'http://{self.host}:{self.port}/file/upload',
                    json={
                        'file_name': os.path.basename(file_path),
                        'total_chunks': total_chunks,
                        'file_hash': file_hash
                    }
                )
                init_response.encoding = 'utf-8'
                if init_response.status_code != 200:
                    try:
                        error_info = init_response.json()
                    except:
                        error_info = f"无法解析响应内容: {init_response.text}"
                    raise Exception(f'分块上传初始化失败: {error_info}')

                with open(file_path, 'rb') as file:
                    for chunk_index in range(total_chunks):
                        chunk = file.read(self.chunk_size)
                        chunk_hash = self.calculate_file_hash(chunk, "sha1")
                        print(f"客户端计算的分块 {chunk_index} 哈希: {chunk_hash}")

                        response = requests.post(
                            f'http://{self.host}:{self.port}/file/chunk_upload/{file_hash}',
                            files={'chunk': chunk},
                            data={'chunk_number': chunk_index, 'chunk_hash': chunk_hash}
                        )
                        response.encoding = 'utf-8'
                        if response.status_code != 200:
                            try:
                                error_info = response.json()
                            except:
                                error_info = f"无法解析响应内容: {response.text}"
                            raise Exception(f'分块上传失败: {error_info}')

                # 最终确认上传
                response = requests.post(f'http://{self.host}:{self.port}/file/complete_upload/{file_hash}')
                response.encoding = 'utf-8'
            else:
                with open(file_path, 'rb') as file:
                    files = {'file': file}
                    response = requests.post(self.url, files=files)
                    response.encoding = 'utf-8'

            if response.status_code == 200:
                print('文件上传成功:', response.json())
                qr_base64 = response.json().get("qr")
                if qr_base64:
                    qr_bytes = base64.b64decode(qr_base64)
                    qr_file_path = os.path.join(qr_save_dir, f"{response.json()['hash']}.png")
                    with open(qr_file_path, 'wb') as qr_file:
                        qr_file.write(qr_bytes)
                    print(f"二维码已保存到 {qr_file_path}")
            else:
                try:
                    error_info = response.json()
                except:
                    error_info = f"无法解析响应内容: {response.text}"
                print('文件上传失败:', error_info)

        except Exception as e:
            print(f"发生未知错误: {e}")